package com.gzu.azy.consumerWeighted.config;

import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;
import java.util.*;

public class WeightedLoadBalancer implements ReactorServiceInstanceLoadBalancer {
    private static final String DEFAULT_WEIGHT = "1";
    private final ObjectProvider<ServiceInstanceListSupplier> supplierProvider;

    public WeightedLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> supplierProvider) {
        this.supplierProvider = supplierProvider;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = supplierProvider.getIfAvailable();
        return supplier.get().next().map(this::getWeightedResponse);
    }

    private Response<ServiceInstance> getWeightedResponse(List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            return new EmptyResponse();
        }
        HashMap<String, Integer> weightMap = new HashMap<>();
        int totalWeight = 0;
        try {
            for (ServiceInstance instance : instances) {
                int weight = Integer.parseInt(instance.getMetadata().getOrDefault("weight", DEFAULT_WEIGHT));
                weightMap.put(instance.getInstanceId(), weight);
                totalWeight += weight;
            }
        } catch (NumberFormatException e) {
            System.err.println(e);
            return new DefaultResponse(instances.get(0));
        }

        int cumulativeWeight = 0;
        int randomValue = new Random().nextInt(totalWeight);
        for (ServiceInstance instance : instances) {
            cumulativeWeight += weightMap.get(instance.getInstanceId());
            if (randomValue < cumulativeWeight) {
                return new DefaultResponse(instance);
            }
        }
        return new DefaultResponse(instances.get(new Random().nextInt(instances.size())));
    }
}